import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool(x)
        x = F.relu(self.conv2(x))
        return x
self.fc1 = nn.Linear(64 * 14 * 14, 128)
self.fc2 = nn.Linear(128, 10)
criterion = nn.CrossEntropyLoss()
def train(model, device, dataloader, optimizer, epoch):
    model.train()
    for batch_idx, (data, target) in enumerate(dataloader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()
def test(model, device, dataloader):
    model.eval()
    total_correct = 0
    with torch.no_grad():
        for data, target in dataloader:
            data, target = data.to(device), target.to(device)
            pred = model(data).argmax(dim=1)
            total_correct += pred.eq(target).sum().item()
    return total_correct / len(dataloader.dataset)
    # 初始化模型、设备、优化器等
model = ConvNet().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# ...其他设置...
for epoch in range(num_epochs):
    train(model, device, train_dataloader, optimizer, epoch)
    accuracy = test(model, device, val_dataloader)
    print(f"Epoch {epoch+1}, Validation Accuracy: {accuracy}")
